import os
import datetime
import pandas as pd
from openpyxl import Workbook
from openpyxl.styles import Alignment, Font, PatternFill, Border, Side
from openpyxl.utils import get_column_letter


def format_survey_files(
    input_dir="input",
    output_dir="output",
    title_text="水表普查清单",
    start_index=0,
    end_index=None,
    columns_to_keep=None,
    save_func=None,
):
    """
    将 input_dir 中所有 Excel 文件按要求格式化并输出到 output_dir。

    参数:
    - title_text: 标题文本，也用于文件名
    - start_index: 开始行号（从0开始）
    - end_index: 结束行号（包含），None表示到最后一行
    - columns_to_keep: 需要保留的列名列表，None表示保留所有列
    """
    os.makedirs(output_dir, exist_ok=True)
    files = [f for f in os.listdir(input_dir) if f.endswith((".xlsx", ".xls"))]
    if not files:
        raise ValueError("未找到 Excel 文件，请重新上传！")

    # 默认删除的列（如果columns_to_keep为None时使用）
    default_deleted_columns = [
        "营业区域",
        "供水片区",
        "抄表周期",
        "口径",
        "抄表员",
        "抄表册",
        "开户日期",
        "单位用户",
        "账户余额",
        "待抄次数",
        "应抄次数",
        "开户人",
        "安装日期",
        "抄表序号",
        "联系人",
        "证件号码",
        "水表属性",
        "供水站",
    ]

    out_files = []

    for file in files:
        path = os.path.join(input_dir, file)
        # 以 None header 读取全部原始行
        df_raw = pd.read_excel(path, header=None, dtype=object)
        if df_raw is None or df_raw.shape[0] == 0:
            continue

        # 把 NaN 转为空字符串以便判断空行
        df_raw = df_raw.fillna("")

        # 如果第二行（索引1）全为空则删除
        if (
            df_raw.shape[0] > 1
            and df_raw.iloc[1].astype(str).map(lambda s: str(s).strip()).eq("").all()
        ):
            df_raw = df_raw.drop(index=1).reset_index(drop=True)

        # 找到包含"序号"作为表头的行（第一处出现）
        header_idx = None
        for i in range(min(len(df_raw), 20)):
            row_vals = df_raw.iloc[i].astype(str)
            if row_vals.str.contains("序号").any():
                header_idx = i
                break
        if header_idx is None:
            raise ValueError(f"{file} 中未找到包含 '序号' 的表头行")

        # 构造表头与数据
        header_row = df_raw.iloc[header_idx].astype(str).tolist()
        data_df = df_raw.drop(index=range(0, header_idx + 1)).reset_index(drop=True)

        if data_df.empty:
            # 仍然输出空表但包含标题与表头
            filtered_df = pd.DataFrame(columns=header_row)
        else:
            data_df.columns = header_row
            # 根据行号筛选数据
            if end_index is None:
                end_index = len(data_df) - 1
            filtered_df = data_df.iloc[start_index : end_index + 1].copy()
            # 重置索引保持连续
            filtered_df = filtered_df.reset_index(drop=True)

        # 列处理
        if columns_to_keep:
            # 只保留指定的列
            available_columns = [
                col for col in columns_to_keep if col in filtered_df.columns
            ]
            missing_columns = [
                col for col in columns_to_keep if col not in filtered_df.columns
            ]

            if missing_columns:
                print(f"警告: 以下列不存在: {missing_columns}")

            if available_columns:
                filtered_df = filtered_df[available_columns]
            else:
                print("警告: 没有找到任何指定的列，将保留所有列")
        else:
            # 删除默认列
            for c in default_deleted_columns:
                if c in filtered_df.columns:
                    filtered_df = filtered_df.drop(columns=[c])

        # 在最后新增备注列
        if "备注" not in filtered_df.columns:
            filtered_df["备注"] = ""

        # 保持原表头顺序
        final_cols = list(filtered_df.columns)

        # 构造最终标题
        now = datetime.datetime.now()
        if start_index == 0 and end_index == len(data_df) - 1:
            range_text = "全部"
        else:
            range_text = f"{start_index+1}-{end_index+1}"

        final_title = f"{now.year}年{now.month}月{title_text}（{range_text}）"

        # 使用 openpyxl 写入并设置格式
        wb = Workbook()
        ws = wb.active

        # 写标题到 A1，并合并单元格、设置样式
        ws.cell(row=1, column=1, value=final_title)
        last_col_idx = len(final_cols) if len(final_cols) > 0 else 1
        ws.merge_cells(start_row=1, start_column=1, end_row=1, end_column=last_col_idx)
        title_cell = ws.cell(row=1, column=1)
        title_cell.font = Font(bold=True, size=18, name="微软雅黑")
        title_cell.alignment = Alignment(
            horizontal="center", vertical="center", wrap_text=True
        )
        title_cell.fill = PatternFill(fill_type="solid", fgColor="D9D9D9")

        # 写表头到第2行
        for ci, col_name in enumerate(final_cols, start=1):
            cell = ws.cell(row=2, column=ci, value=col_name)
            cell.font = Font(bold=True)
            cell.alignment = Alignment(
                horizontal="center", vertical="center", wrap_text=True
            )

        # 写数据从第3行开始
        for r_idx, row in enumerate(
            filtered_df.itertuples(index=False, name=None), start=3
        ):
            for c_idx, value in enumerate(row, start=1):
                cell = ws.cell(row=r_idx, column=c_idx, value=value)
                cell.alignment = Alignment(
                    horizontal="center", vertical="center", wrap_text=True
                )

        # 设置行高
        ws.row_dimensions[1].height = 40
        for r in range(2, ws.max_row + 1):
            ws.row_dimensions[r].height = 25

        # 为单元格添加细边框
        thin = Side(border_style="thin", color="000000")
        cell_border = Border(left=thin, right=thin, top=thin, bottom=thin)
        for c in range(1, last_col_idx + 1):
            for r in range(1, ws.max_row + 1):
                ws.cell(row=r, column=c).border = cell_border

        # 自适应列宽
        for idx in range(1, len(final_cols) + 1):
            letter = get_column_letter(idx)
            col_name = final_cols[idx - 1]
            try:
                # 计算列宽：表头宽度和内容最大宽度的最大值
                header_width = len(str(col_name))
                content_width = filtered_df[col_name].astype(str).str.len().max()
                max_width = max(header_width, content_width)
                adjusted_width = min(max_width + 2, 50)  # 限制最大宽度
                ws.column_dimensions[letter].width = adjusted_width
            except Exception:
                ws.column_dimensions[letter].width = 15

        # 输出文件名
        safe_name = "".join(ch for ch in final_title if ch not in r'\/:*?"<>|').strip()
        out_path = os.path.join(output_dir, f"{safe_name}.xlsx")

        # 保存
        if save_func:
            save_func(wb, out_path)
        else:
            wb.save(out_path)
        out_files.append(out_path)

    return out_files
